import pandas as pd

from tqdm import tqdm
from datetime import datetime, timedelta

from pyecharts import options as opts
from pyecharts.charts import Bar, Tab, Line, Page, Grid, Timeline
from pyecharts.globals import ThemeType


tjd_legend_opts = opts.LegendOpts(pos_top="7%", is_show=True, textstyle_opts=opts.TextStyleOpts(font_size=18))
tjd_text_opts = opts.TextStyleOpts(font_size=18)

def multi_macro_bar_or_line(df, chart_type):
    """
    绘制多条线的柱形图或线图
    """
    
    if chart_type == 'bar':
        b = Bar(init_opts=opts.InitOpts(height='600px', theme=ThemeType.ROMA))
    elif chart_type == 'line':
        b = Line(init_opts=opts.InitOpts(height='600px', theme=ThemeType.ROMA))
    
    b.add_xaxis([s.strftime("%Y-%m-%d") for s in df.index]) # x数据（常是索引）
    
    for y_name in df.columns:
        se = df[y_name]
        b.add_yaxis(y_name, y_axis=se.to_list()) # y数据
    
    b.set_global_opts(
        datazoom_opts=[
            opts.DataZoomOpts(range_start=80, range_end=100),
            opts.DataZoomOpts(type_="inside")],
        toolbox_opts=opts.ToolboxOpts(),
        yaxis_opts=opts.AxisOpts(name="%", axislabel_opts=opts.LabelOpts(font_size=16)),
        xaxis_opts=opts.AxisOpts(name="日期", axislabel_opts=opts.LabelOpts(font_size=16)),
        legend_opts=tjd_legend_opts
        )
    return b


def multi_grid_line(df):
    """
    绘制四方格线型图
    """
    df = df.dropna(how='all')
    df = df.round(2)
    grid = Grid(init_opts=opts.InitOpts(height="600px"))
    assert df.shape[1] <= 4
    aggs = [
        {'pos_right':'54%', 'pos_bottom':'54%'},
        {'pos_left':'54%', 'pos_bottom':'54%'},
        {'pos_right':'54%', 'pos_top':'54%'},
        {'pos_left':'54%', 'pos_top':'54%'},
            ]
    label_pos = [
        {'pos_right':'70%', 'pos_bottom':'94%'},
        {'pos_left':'70%', 'pos_bottom':'94%'},
        {'pos_right':'70%', 'pos_bottom':'46%'},
        {'pos_left':'70%', 'pos_bottom':'46%'},
        ]
    for i, y_name in enumerate(df.columns.to_list()):
        se = df[y_name]
        se = se.dropna()
        se = se[-100:]
        b = Line()
        b.add_xaxis(se.index.to_list()) # x数据（常是索引）
        b.add_yaxis(y_name, y_axis=se.to_list()) # y数据
        b.set_global_opts(
            title_opts=opts.TitleOpts(title=se.name, **aggs[i]),
            toolbox_opts=opts.ToolboxOpts(),
            xaxis_opts=opts.AxisOpts(name="date"),
            legend_opts=opts.LegendOpts(**label_pos[i], is_show=True, textstyle_opts=tjd_text_opts)
            )
        b.set_series_opts(label_opts=opts.LabelOpts(is_show=False))
        print(i,aggs[i])
        grid.add(b, grid_opts=opts.GridOpts(**aggs[i]))
    return grid


def year_timeline_seasonal(df):
    # 轮播图，按年播放
    cols = df.columns.to_list()
    
    df = df.dropna(how='all')
    df = df.round(2)
    df.index = pd.to_datetime(df.index)

    years = df.index.year.unique().to_list()
    year_dic = {y_name:to_monthly_mean(df[y_name]) for y_name in cols}
    
    tml = Timeline(init_opts=opts.InitOpts(width='1400px', height='900px'))
    for year in years:
        b = Line()
        b.add_xaxis([f'{i}月' for i in range(1, 13)])
        for y_name in cols:
            seasonal_se = year_dic[y_name][year]
            b.add_yaxis(y_name, seasonal_se.to_list())
            b.set_global_opts(title_opts=opts.TitleOpts(f"{year}年"),
                              legend_opts=tjd_legend_opts
                              )
        tml.add(b, f"{year}年")
    return tml


def item_timeline_seasonal(df, freq, year_dic):
    # 轮播图，按数据播放（展示季节性, 数据为日度或周度或月）

    cols = df.columns.to_list()
    
    df = df.dropna(how='all')
    df = df.round(2)
    df.index = pd.to_datetime(df.index)

    years = df.index.year.unique().to_list()
    
    tml = Timeline(init_opts=opts.InitOpts(width='1400px', height='600px'))
    for y_name in cols:
        b = Line()
        
        if freq == 'D':
            xaxis_series = year_dic[y_name].index.to_list()
            smooth = True
        elif freq == 'M':
            xaxis_series = [f'{i}月' for i in range(1, 13)]
            smooth = False
        b.add_xaxis(xaxis_series)
        
        for year in years:
            seasonal_se = year_dic[y_name][year]
            if isinstance(seasonal_se, pd.DataFrame):
                seasonal_se = seasonal_se.iloc[:,0]
            if len(seasonal_se.dropna()) >= 1:
                if year == 2024:
                    b.add_yaxis(
                        f"{year}年", seasonal_se.to_list(),
                        linestyle_opts=opts.LineStyleOpts(width=2, color='black'),
                        itemstyle_opts=opts.ItemStyleOpts(color='black'),  # 设置数据点颜色
                        is_connect_nones=True, is_smooth=smooth,
                        label_opts=(opts.LabelOpts(is_show=False))
                        )
                else:
                    b.add_yaxis(f"{year}年", seasonal_se.to_list(),
                                is_connect_nones=True, is_smooth=smooth,
                                label_opts=(opts.LabelOpts(is_show=False)))
                
                b.set_global_opts(
                    title_opts=opts.TitleOpts(y_name),
                    legend_opts=opts.LegendOpts(pos_right="3%", pos_top="12%",
                                                orient='vertical',
                                                textstyle_opts=tjd_text_opts),
                    toolbox_opts=opts.ToolboxOpts(),
                                  )
            
        tml.add(b, y_name)
    return tml


def xy_match_same_freq_plot(df_cot, idx_fill, tztd_ratio):
    cols = df_cot.columns.to_list()
    b = Line()
    idx = [str(s)[:10] for s in df_cot.index]
    b.add_xaxis(idx)
    b.add_yaxis(cols[0], df_cot[cols[0]].to_list(), yaxis_index=0,
                color='black') # 添加第一条曲线，默认使用左侧y轴
    # 扩展一个右侧的y轴
    b.extend_axis(
        yaxis=opts.AxisOpts(name="特征X", type_="value", position="right")
    )
    b.add_yaxis(cols[1], df_cot[cols[1]], yaxis_index=1, color='red') # 添加第二条曲线，指定使用右侧y轴
    b.set_series_opts(
        markarea_opts=opts.MarkAreaOpts(
            data=[opts.MarkAreaItem(x=s) for s in idx_fill],
            itemstyle_opts=opts.ItemStyleOpts(color='pink',opacity=0.1),
        ),
                      )
    # 全局配置项
    b.set_global_opts(
        datazoom_opts=[
            opts.DataZoomOpts(range_start=10, range_end=100),
            opts.DataZoomOpts(type_="inside")],
        title_opts=opts.TitleOpts(title=f"因子关联分析>同涨同跌率:{100*tztd_ratio:.2f}%"),
        yaxis_opts=opts.AxisOpts(name="指标Y"),
        toolbox_opts=opts.ToolboxOpts(),
        tooltip_opts=opts.TooltipOpts(trigger="axis", axis_pointer_type="cross"),
        legend_opts=opts.LegendOpts(pos_top="6%", is_show=True)
    )
    return b